{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Random Number Generation for JIT Compilation\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/brainpy/brainpy/blob/master/docs_version2/tutorial_math/random_number_generation.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/brainpy/brainpy/blob/master/docs_version2/tutorial_math/random_number_generation.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "[Chaoming Wang](mailto:chao.brain@qq.com)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Although ``brainpy.math.random`` is designed to be seamlessly compatible with ``numpy.random``, there are still some differences under the context of JIT compilation.\n",
    "\n",
    "In this section, we are going to talk about how to program a JIT-compatible code with ``brainpy.math.random``."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:03.969922Z",
     "start_time": "2025-10-06T04:03:03.966561Z"
    }
   },
   "source": [
    "import brainpy as bp\n",
    "import brainpy.math as bm\n",
    "import numpy as np\n",
    "\n",
    "# bm.set_platform('cpu')"
   ],
   "outputs": [],
   "execution_count": 28
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:03.989006Z",
     "start_time": "2025-10-06T04:03:03.982761Z"
    }
   },
   "source": [
    "bp.__version__"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'3.0.0'"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 29
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Using ``bm.random`` outside functions to JIT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Using ``brainpy.math.random`` outside of functions to JIT is the same as using ``numpy.random``.\n",
    "\n",
    "This usage corresponds to the cases that generating random data for further processing. For example,"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.018865Z",
     "start_time": "2025-10-06T04:03:04.013080Z"
    }
   },
   "source": [
    "bm.random.rand(10)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.3089733 , 0.89542127, 0.02714849, 0.94584775, 0.77089345,\n",
       "       0.32520366, 0.69955456, 0.37229824, 0.43667316, 0.02874231],      dtype=float32)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 30
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.038321Z",
     "start_time": "2025-10-06T04:03:04.031639Z"
    }
   },
   "source": [
    "np.random.rand(10)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.18021362, 0.57020223, 0.91773743, 0.32111789, 0.592986  ,\n",
       "       0.32347693, 0.66178114, 0.11078982, 0.20402456, 0.80672322])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 31
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "When you are using API functions in ``brainpy.math.random``, actually you are calling functions in a default ``RandomState``. Similarly, ``numpy.random`` also has a default ``RandomState``. Calling a random function in ``numpy.random`` module corresponds to calling the random function in this default NumPy ``RandomState``."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.063855Z",
     "start_time": "2025-10-06T04:03:04.059525Z"
    }
   },
   "source": [
    "bm.random.DEFAULT"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RandomState([2354496134 1053393286])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 32
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Using ``bm.random`` inside a function to JIT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "If you are using random sampling in a JIT function, **there are things you need to pay attention to**. Otherwise, the error is likely to raise."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "As I have stated above, ``brainpy.math.random`` functions are using the default ``RandomState``. A ``RandomState`` is an instance of brainpy [Variable](./arrays_and_variables.ipynb), denoting that it has values to change after calling any its built-in random function. What's changing is the `key` of a ``RandomState``. For instance,"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.086112Z",
     "start_time": "2025-10-06T04:03:04.081634Z"
    }
   },
   "source": [
    "bm.random.rand(1)\n",
    "print('Now, the DEFAULT is', bm.random.DEFAULT)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Now, the DEFAULT is RandomState([3949348841  188664647])\n"
     ]
    }
   ],
   "execution_count": 33
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.114951Z",
     "start_time": "2025-10-06T04:03:04.109470Z"
    }
   },
   "source": [
    "bm.random.rand(1)\n",
    "print('Now, the DEFAULT is', bm.random.DEFAULT)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Now, the DEFAULT is RandomState([1197037675  819676900])\n"
     ]
    }
   ],
   "execution_count": 34
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Therefore, if you do not specify this DEFAULT ``RandomState`` you are using, repeatedly calling random functions in ``brainpy.math.random`` module will not get what you want, because its `key` cannot be updated. For instance,"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.134587Z",
     "start_time": "2025-10-06T04:03:04.127958Z"
    }
   },
   "source": [
    "@bm.jit\n",
    "def get_data():\n",
    "    return bm.random.random(2)"
   ],
   "outputs": [],
   "execution_count": 35
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.251134Z",
     "start_time": "2025-10-06T04:03:04.161623Z"
    }
   },
   "source": [
    "get_data()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.6563661 , 0.82885206], dtype=float32)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 36
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.264344Z",
     "start_time": "2025-10-06T04:03:04.258890Z"
    }
   },
   "source": [
    "get_data()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.8161644, 0.5287926], dtype=float32)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 37
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "A correct way is explicitly declaring you are using this DEFAULT variable in the JIT transformation."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.280423Z",
     "start_time": "2025-10-06T04:03:04.276196Z"
    }
   },
   "source": [
    "bm.random.seed()"
   ],
   "outputs": [],
   "execution_count": 38
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.298031Z",
     "start_time": "2025-10-06T04:03:04.292460Z"
    }
   },
   "source": [
    "from functools import partial\n",
    "\n",
    "@partial(bm.jit)\n",
    "def get_data_v2():\n",
    "    return bm.random.random(2)"
   ],
   "outputs": [],
   "execution_count": 39
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.390467Z",
     "start_time": "2025-10-06T04:03:04.304539Z"
    }
   },
   "source": [
    "get_data_v2()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.18351912, 0.07486355], dtype=float32)"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 40
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.404583Z",
     "start_time": "2025-10-06T04:03:04.399766Z"
    }
   },
   "source": [
    "get_data_v2()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.6289803, 0.797102 ], dtype=float32)"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 41
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Or, declare the function as a `BrainPyObject`, then use ``jit()``."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.421710Z",
     "start_time": "2025-10-06T04:03:04.417679Z"
    }
   },
   "source": [
    "@bm.jit\n",
    "def get_data_v3():\n",
    "    return bm.random.random(2)"
   ],
   "outputs": [],
   "execution_count": 42
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.531244Z",
     "start_time": "2025-10-06T04:03:04.435887Z"
    }
   },
   "source": [
    "get_data_v3()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.64915967, 0.37961698], dtype=float32)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 43
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.545961Z",
     "start_time": "2025-10-06T04:03:04.540891Z"
    }
   },
   "source": [
    "get_data_v3()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.28215742, 0.5177479 ], dtype=float32)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 44
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Using ``RandomState`` for objects to JIT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Another way I recommend is using instances of ``RandomState`` for objects to JIT. For example, you can initialize a ``RandomState`` in the ``__init__()`` function, then using the initialized ``RandomState`` anywhere."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.561633Z",
     "start_time": "2025-10-06T04:03:04.557509Z"
    }
   },
   "source": [
    "class MyOb(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.rng = bm.random.RandomState(123)\n",
    "\n",
    "    def __call__(self):\n",
    "        size = (50, 100)\n",
    "        u = self.rng.random(size)\n",
    "        v = self.rng.uniform(size=size)\n",
    "        z = bm.sqrt(-2 * bm.log(u)) * bm.cos(2 * bm.pi * v)\n",
    "        return z"
   ],
   "outputs": [],
   "execution_count": 45
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.577406Z",
     "start_time": "2025-10-06T04:03:04.573250Z"
    }
   },
   "source": [
    "ob = bm.jit(MyOb())"
   ],
   "outputs": [],
   "execution_count": 46
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:04.863417Z",
     "start_time": "2025-10-06T04:03:04.586251Z"
    }
   },
   "source": [
    "ob()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 1.714168  , -0.93755597, -1.566088  , ..., -2.300819  ,\n",
       "        -0.3244472 , -1.0213778 ],\n",
       "       [ 0.22194532,  0.86869   , -1.8543271 , ...,  0.6031633 ,\n",
       "         0.5587094 ,  0.17947507],\n",
       "       [-0.04175869, -0.25695515,  0.41855374, ...,  0.7097617 ,\n",
       "        -1.0369484 , -2.6332123 ],\n",
       "       ...,\n",
       "       [-0.5513834 , -1.2007083 , -0.21716705, ...,  2.4005718 ,\n",
       "        -1.97702   , -1.9229662 ],\n",
       "       [ 1.9620371 , -1.030306  ,  0.5033778 , ...,  1.0114168 ,\n",
       "         0.5670819 ,  0.3017683 ],\n",
       "       [-1.576822  , -1.2365927 ,  0.87836343, ..., -1.5918341 ,\n",
       "         0.7913105 , -1.8031495 ]], dtype=float32)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 47
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Note that any ``Variable`` instance which can be directly accessed by ``self.`` is able to be automatically found by brainpy's JIT transformation functions.\n",
    "Therefore, in this case, we do not need to pass the ``rng`` into the `dyn_vars` in ``bm.jit()`` function."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
